package com.chukun.flink.stream.state.key;

import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.SourceFunction;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.concurrent.TimeUnit;

/**
 * @author chukun
 * @version 1.0.0
 * @description keyState基本操作
 * @createTime 2022年05月15日 22:39:00
 */
public class KeyStateBase {
    /**
     * 自定义非并行数据源，该数据源操作符的并行度为 1
     */
    public static class StateSource implements SourceFunction<Tuple2<Integer, Integer>> {

        private static final Logger logger = LoggerFactory.getLogger(StateSource.class);

        private int  counter = 0;

        @Override
        public void run(SourceContext<Tuple2<Integer, Integer>> sourceContext) throws Exception {
            while (true) {
                // 向数据流中发送元素， f0为 counter%5 的数， f1为counter的值
                sourceContext.collect(new Tuple2<>(counter % 5, counter));
                logger.info("send data ; {} , {} ", counter % 5, counter);
                counter++;
                TimeUnit.SECONDS.sleep(1);
            }
        }

        @Override
        public void cancel() {}
    }

    /**
     * 封装了创建数据源函数和对数据源函数生成的数据流进行分组操作，返回keyedStream
     * @param env
     * @return
     */
    public static KeyedStream<Tuple2<Integer,Integer>, Integer> before(StreamExecutionEnvironment env) {
        // 设置作业的全局并行度
        env.setParallelism(2);
        //定义数据源，每1s发送一个元素
        DataStream<Tuple2<Integer, Integer>> inputStream = env.addSource(new StateSource());
        // 将数据流中的元素索引为0的字段作为key去分流
        KeyedStream<Tuple2<Integer, Integer>, Integer> keyedStream = inputStream.keyBy(new KeySelector<Tuple2<Integer, Integer>, Integer>() {
            @Override
            public Integer getKey(Tuple2<Integer, Integer> tuple2) throws Exception {
                return tuple2.f0;
            }
        });
        return keyedStream;
    }
}
